#include "Storage.h"

#include <Support\Log.h>
#include <Support\GAssert.h>

#ifndef NULL
#define NULL nullptr
#endif

StorageManager StorageManager::theInstance;

StorageManager& StorageManager::getInstance()
{
	return theInstance;
}

U8* StorageManager::allocate(size_t size_bytes, size_t alignment)
{	
	ASSERT(isInitialized == true, "Storage Manager not initialized");
	ASSERT((size_bytes != 0 && alignment >= 1 && alignment <= 128 != 0),
		"Invalid Memory Request, Size = %u, Alignment = %u", 
		size_bytes, alignment);
	ASSERT((alignment & (alignment - 1)) == 0, "Alignment must be a power of 2");

	Allocator* prevFreeBlock = NULL;
	Allocator* freeBlock = freeStorage;

	size_t expandedSize_bytes = sizeof(Allocator) + alignment + size_bytes;

	while (NULL != freeBlock)
	{
		if (expandedSize_bytes < freeBlock->length)
			break;
		prevFreeBlock = freeBlock;
		freeBlock = freeBlock->next;
	}
	ASSERT(expandedSize_bytes < freeBlock->length,
		"Not enough storage to allocate %u bytes",
		expandedSize_bytes);

	//If allocations in the remaining memory will be impossible
	if (freeBlock->length - expandedSize_bytes <= sizeof(Allocator))
	{
		//Increase allocation size instead of creating a new FreeBlock
		expandedSize_bytes = freeBlock->length;

		if (prevFreeBlock != NULL)
			prevFreeBlock->next = freeBlock->next;
		else
			freeStorage = freeBlock->next;
	}
	else
	{
		uintptr_t nextFreeStorageBlock = reinterpret_cast<uintptr_t>(freeBlock)+expandedSize_bytes;
		Allocator* next_block = reinterpret_cast<Allocator* >(nextFreeStorageBlock);
		next_block->length = freeBlock->length - expandedSize_bytes;
		next_block->next = freeBlock->next;
		freeBlock->next = NULL;


		if (prevFreeBlock != NULL)
			prevFreeBlock->next = next_block;
		else
			freeStorage = next_block;
	}

	uintptr_t rawAddress = reinterpret_cast<uintptr_t>(freeBlock);

	size_t mask = (alignment - 1);
	uintptr_t rawDataAddress = rawAddress + sizeof(Allocator);
	uintptr_t misalignment = (rawDataAddress & mask);
	ptrdiff_t adjustment = alignment - misalignment;
	ASSERT(adjustment < 256, "Internal Error, Adjustment out of boundary");
	uintptr_t alignedAddress = rawDataAddress + adjustment;
	U8 * pAlignedMem = reinterpret_cast<U8* >(alignedAddress);
	pAlignedMem[-1] = static_cast<U8>(adjustment);

	freeBlock->next = usedStorage;
	freeBlock->length = expandedSize_bytes;
	usedStorage = freeBlock;

	return static_cast<U8* >(pAlignedMem);
}

void StorageManager::deallocate(void* pMem)
{
	if (!isInitialized)
		return;
	const U8 * pAlignedMem
		= reinterpret_cast<const U8 * >(pMem);
	uintptr_t alignedAddress
		= reinterpret_cast<uintptr_t>(pMem);
	ptrdiff_t adjustment
		= static_cast<ptrdiff_t>(pAlignedMem[-1]);

	uintptr_t block_start = alignedAddress - adjustment - sizeof(Allocator);
	Allocator* pRawMem = reinterpret_cast<Allocator* >(block_start);

	Allocator* prevAlloc = NULL;
	Allocator* alloc;
	for (alloc = usedStorage;
		alloc != NULL;
		prevAlloc = alloc, alloc = alloc->next)
	{
		if (alloc == pRawMem)
		{
			if (NULL == prevAlloc)
				usedStorage = alloc->next;
			else
				prevAlloc->next = alloc->next;
			break;
		}
	}
	prevAlloc = NULL;
	bool mergedBefore = false;
	bool mergedAfter = false;
	alloc = freeStorage;
	Allocator* nextpRawMem;

	while (alloc != NULL)
	{
		nextpRawMem = mergeAllocators(*pRawMem, *alloc);
		if (NULL == nextpRawMem)
		{
			prevAlloc = alloc;
			alloc = alloc->next;
		}
		else
		{
			pRawMem = nextpRawMem;
			alloc = freeStorage;
		}
	}
	pRawMem->next = freeStorage;
	freeStorage = pRawMem;
}

U8* StorageManager::getStorage(size_t& storageLength)
{
	size_t numBytesxUnit[] = { 1, 1, 1, 1 };
	uint MAX_NUM_UNITS = sizeof(numBytesxUnit) / sizeof(*numBytesxUnit);
	uint currentUnitIndex = 0;

	static const uint amount[] = { 512, 1 };
	uint MAX_NUM_AMOUNT = sizeof(amount) / sizeof(*amount);
	uint currentAmountIndex = 0;

	for (int i = MAX_NUM_UNITS - 2; i >= 0; i--)
		numBytesxUnit[i] = numBytesxUnit[i + 1] * 1024;

	U8*	storageAquired = NULL;

	while (NULL == storageAquired && currentUnitIndex < MAX_NUM_UNITS)
	{
		currentAmountIndex = 0;
		while (NULL == storageAquired && currentUnitIndex < MAX_NUM_UNITS && currentAmountIndex < MAX_NUM_AMOUNT)
		{
			storageLength = (size_t)amount[currentAmountIndex] * numBytesxUnit[currentUnitIndex];
			if (storageLength)
			{
				storageAquired = reinterpret_cast<U8* >(malloc(storageLength));
				if (NULL != storageAquired)
					break;
			}
			currentAmountIndex++;
		}
		if (NULL != storageAquired)
			break;
		currentUnitIndex++;
	}
	return storageAquired;
} 

void StorageManager::dequeueFreeStorage(Allocator& allocator)
{
	if (&allocator == freeStorage)
	{
		freeStorage = allocator.next;
		allocator.next = NULL;
	}
	else
		for (Allocator* alloc = freeStorage; NULL != alloc; alloc = alloc->next)
		{
			if (&allocator != alloc->next)
				continue;
			alloc->next = allocator.next;
			allocator.next = NULL;
		}	
}

StorageManager::Allocator* StorageManager::mergeAllocators(Allocator& first, Allocator& second)
{
	const U8 * pFirst
		= reinterpret_cast<const U8 * >(&first);
	const U8 * pSecond
		= reinterpret_cast<const U8 * >(&second);
	const U8* nextToFirst = pFirst + first.length;
	const U8* nextToSecond = pSecond + second.length;

	if (nextToFirst != pSecond && nextToSecond != pFirst)
		return NULL;

	dequeueFreeStorage(first);
	dequeueFreeStorage(second);

	if (nextToFirst == pSecond)
	{
		first.next = second.next;
		first.length += second.length;
		return &first;
	}
	second.next = first.next;
	second.length += first.length;
	return &second;
}

bool StorageManager::initialize()
{
	if (isInitialized)
		return true;

	freeStorage = NULL;
	usedStorage = NULL;

	for (uint i = 0; i < MAX_STORAGE_POOL; i++)
	{
		size_t storageLength = 0;
		Allocator* storage = reinterpret_cast<Allocator*> (getStorage(storageLength));
		if (NULL != storage)
		{
			pool[i] = storage;
			storage->next = freeStorage;
			storage->length = storageLength;
			freeStorage = storage;
		}
		else pool[i] = NULL;
	}

	isInitialized = (freeStorage) ? true : false;
	return isInitialized;
}

bool StorageManager::shutdown()
{
	if (! isInitialized)
		return true;

	for (uint i = 0; i < MAX_STORAGE_POOL; i++)
		if (NULL != pool[i])
			free(pool[i]);

	isInitialized = false;
	return true;
}

void* StorageManager::New(size_t size_bytes, size_t alignment)
{
	return allocate(size_bytes, alignment);
}

void StorageManager::Delete(void * pMem)
{
	deallocate(pMem);
}